import os
import hashlib
import argparse
from tqdm import tqdm


def calculate_file_hash(file_path):
    """计算并返回给定文件的SHA256哈希值"""
    hasher = hashlib.sha256()
    with open(file_path, 'rb') as f:
        buf = f.read()
        hasher.update(buf)
    return hasher.hexdigest()


def create_mirror_structure(src_dir, dst_dir):
    """
    递归地扫描src_dir并在dst_dir创建相同的文件夹结构。
    对于每个文件，在相应的dst_dir位置创建一个txt文件，内容为原文件的哈希值，并保留原始文件的扩展名。
    """
    # 计算总文件数，用于进度条
    total_files = sum(len(files) for _, _, files in os.walk(src_dir))
    
    with tqdm(total=total_files, desc='Processing files', unit='file') as pbar:
        for root, dirs, files in os.walk(src_dir):
            relative_path = os.path.relpath(root, src_dir)
            dst_root = os.path.join(dst_dir, relative_path)

            if not os.path.exists(dst_root):
                os.makedirs(dst_root)

            for file_name in files:
                src_file_path = os.path.join(root, file_name)
                hash_content = calculate_file_hash(src_file_path)
                
                # 保留原始文件的扩展名
                dst_file_name = f"{file_name}.txt"
                dst_file_path = os.path.join(dst_root, dst_file_name)

                with open(dst_file_path, 'w') as txt_file:
                    txt_file.write(hash_content)
                
                pbar.update(1)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate a directory structure based on file hashes.')
    parser.add_argument('--scan-dir', required=True, help='The directory to scan')
    parser.add_argument('--hash-dir', required=True, help='The directory to save hash structures')

    args = parser.parse_args()

    if not os.path.isdir(args.scan_dir):
        print(f"Error: {args.scan_dir} is not a valid directory.")
    else:
        create_mirror_structure(args.scan_dir, args.hash_dir)